﻿//////////////////////////////////////////////////////////////////////////
// @brief 二叉堆实现
// @author youdang0809
// @date 2012-10-19
//////////////////////////////////////////////////////////////////////////

#include "heap.h"
#include <stdlib.h>
#include <stdio.h>

//////////////////////////////////////////////////////////////////////////
// 返回节点的父节点编号
// @param index 当前节点的编号
// @return 父节点编号，若当前节点为根节点，则返回-1
//////////////////////////////////////////////////////////////////////////
static int _heap_parent_index(int index)
{
    return index ? (index - 1) >> 1 : -1;
}

//////////////////////////////////////////////////////////////////////////
// 返回节点的左子节点编号
// @param index 当前节点的编号
// @return 左子节点编号
// @attention 调用之前需判断index是否有左子节点
//////////////////////////////////////////////////////////////////////////
static int _heap_left_index(int index)
{
    return (index << 1) + 1;
}

//////////////////////////////////////////////////////////////////////////
// 返回节点的右子节点编号
// @param index 当前节点的编号
// @return 右子节点编号
// @attention 调用之前需判断index是否有右子节点
//////////////////////////////////////////////////////////////////////////
static int _heap_right_index(int index)
{
    return (index << 1) + 2;
}

//////////////////////////////////////////////////////////////////////////
// 交换两个节点的内容
// @param a a节点地址
// @param b b节点地址
//////////////////////////////////////////////////////////////////////////
static void _heap_node_swap(struct heap_node *a, struct heap_node *b)
{
    void *tmp = a->data;

    a->key ^= b->key;
    b->key ^= a->key;
    a->key ^= b->key;

    a->data = b->data;
    b->data = tmp;
}

//////////////////////////////////////////////////////////////////////////
// 比较两个key的值
// @param hp 指向堆的指针
// @param key1 键值1
// @param key2 键值2
// @return 若是小顶堆，key1小于key2时返回1，否则返回0；若是大顶堆，作相反处理
//////////////////////////////////////////////////////////////////////////
static int _heap_compare(struct heap *hp, int key1, int key2)
{
    if(MAX_HEAP == hp->type)
        return key1 > key2 ? 1 : 0;
    else
        return key1 < key2 ? 1 : 0;
}

//////////////////////////////////////////////////////////////////////////
// 将当前节点上移至适当位置
// @param hp 指向堆的指针
// @param index 当前节点编号
//////////////////////////////////////////////////////////////////////////
static void _heap_up(struct heap *hp, int index)
{
    if(-1 == _heap_parent_index(index))
        return;

    // # 大顶堆：如果节点大于父节点，则交换之，递归向上检查
    // # 小顶堆：如果节点小于父节点，则交换之，递归向上检查
    if(_heap_compare(hp, hp->nodes[index].key, hp->nodes[_heap_parent_index(index)].key))
    {
        _heap_node_swap(&hp->nodes[index], &hp->nodes[_heap_parent_index(index)]);
        _heap_up(hp, _heap_parent_index(index));
    }
}

//////////////////////////////////////////////////////////////////////////
// 将当前节点下移至适当位置
// @param hp 指向堆的指针
// @param index 当前节点编号
//////////////////////////////////////////////////////////////////////////
static void _heap_down(struct heap *hp, int index)
{
    if(index >= hp->size)
        return;
    
    // 对溢出进行判断
    if(index > _heap_parent_index(INT_MAX))
    {
        // (INT_MAX-1)/2 位置上的节点没有右子结点，(INT_MAX-1)/2+1 ~ INT_MAX 位置上的
        // 节点没有子节点
        return;
    }
    else if(index == _heap_parent_index(INT_MAX))
    {
        // 如果当前需要heap_down的节点是最满情况下（共有INT_MAX个节点）最后一个节点的父
        // 节点
        // # 大顶堆：如果当前节点的key大于左子节点的key
        // # 小顶堆：如果当前节点的key小于左子节点的key
        if(_heap_compare(hp, hp->nodes[index].key, hp->nodes[_heap_left_index(index)].key))
            return;
        else
        {
            _heap_node_swap(&hp->nodes[index], &hp->nodes[_heap_left_index(index)]);
            return;
        }
    }
    else
    {
        // 找到键值最大或最小的节点编号
        int endpoint = index;
        if(_heap_left_index(index) < hp->size && _heap_compare(hp, hp->nodes[_heap_left_index(index)].key, hp->nodes[endpoint].key))
            endpoint = _heap_left_index(index);
        if(_heap_right_index(index) < hp->size && _heap_compare(hp, hp->nodes[_heap_right_index(index)].key, hp->nodes[endpoint].key))
            endpoint = _heap_right_index(index);
        if(endpoint != index)
        {
            _heap_node_swap(&hp->nodes[index], &hp->nodes[endpoint]);
            _heap_down(hp, endpoint);
        }
        return;
    }
}

//////////////////////////////////////////////////////////////////////////
// 在index为根的子树中查找节点
// @param hp 指向堆的指针
// @param key 需要查找的节点的key
// @param index 子树的根节点的编号
// @return 若找到，返回节点的地址，否则返回NULL
//////////////////////////////////////////////////////////////////////////
static struct heap_node *_heap_find(struct heap *hp, int key, int index)
{
    // # 小顶堆：若key小于当前节点，则停止查找，否则搜索左右子树
    // # 大顶堆：若key大于当前节点，则停止查找，否则搜索左右子树
    if(index >= hp->size)
        return NULL;
    else if(hp->nodes[index].key == key)
        return &hp->nodes[index];
    else if(_heap_compare(hp, key, hp->nodes[index].key))
        return NULL;
    else
    {
        struct heap_node *ptr = _heap_find(hp, key, _heap_left_index(index));
        if(NULL == ptr)
            ptr = _heap_find(hp, key, _heap_right_index(index));
        return ptr;
    }
}

//////////////////////////////////////////////////////////////////////////
int heap_init(struct heap *hp, int type)
{
    if(NULL == hp)
        return -1;

    hp->capacity = 1;
    hp->size = 0;
    hp->type = type;
    hp->nodes = (struct heap_node *)malloc(sizeof(struct heap_node));
    if(NULL == hp->nodes)
        return -1;
    return 0;
}

int heap_insert(struct heap *hp, int key, void *data)
{
    if(NULL == hp || INT_MAX == hp->size || NULL != heap_find(hp, key))
        return -1;

    // 增大存储空间
    if(hp->size == hp->capacity)
    {
        struct heap_node *tmp = (struct heap_node *)realloc(hp->nodes,
                                                            sizeof(struct heap_node) * (hp->capacity << 1));
        if(NULL == tmp)
            return -1;
        hp->nodes = tmp;
        tmp = NULL;
        hp->capacity <<= 1; // 增大存储空间大小值
    }

    // 将数据插入到最后一个节点上
    hp->nodes[hp->size].key = key;
    hp->nodes[hp->size].data = data;

    // 将节点放置到正确的位置
    _heap_up(hp, hp->size);

    hp->size++;
    return 0;
}

struct heap_node *heap_delete(struct heap *hp)
{
    struct heap_node *root = NULL;
    if(NULL == hp || 0 == hp->size)
        return root;
    
    // 交换根节点和最后一个节点
    _heap_node_swap(&hp->nodes[0], &hp->nodes[hp->size - 1]);
    root = &hp->nodes[hp->size - 1];
    hp->size--;

    // 为新根节点找到合适的位置
    if(hp->size > 0)
        _heap_down(hp, 0);

    return root;
}

struct heap_node *heap_find(struct heap *hp, int key)
{
    if(NULL == hp || 0 == hp->size)
        return NULL;
    return _heap_find(hp, key, 0);
}

void heap_sort(struct heap *hp)
{
    int index = 0, size = 0;
    if(NULL == hp || 0 == hp->size)
        return;

    index = hp->size - 1;
    size = hp->size;        // 保存原来的堆大小
    while(index > 0)
    {
        // 将根节点换到最后位置上
        _heap_node_swap(&hp->nodes[0], &hp->nodes[index]);
        index--;
        hp->size--;         // 去掉最后一个节点
        _heap_down(hp, 0);
    }
    hp->size = size;        // 恢复原来的堆大小
}